#!/usr/bin/python
# -*-coding:utf-8-*-
import os
from time import time

import pandas as pd

### 底层读取数据的依赖（不提供）
from zbc_factor_lib.base.factors_library_base import NewRQFactorLib as DataReader

from base.trade_date_data_generation import get_every_fixed_frequency_trade_date
from base.utils import loadFile
from base.utils import writeToFile

from factor_test.factor_selection_test import FactorTest

# from factor_selection.get_gp_factor import FactorBarraStrippingForTest as UsedFactorBarraStripping
from factor_selection.get_gp_factor import factor_barra_strip_api
from factor_selection.get_gp_factor import get_factor_by_model

from base.KeyTransform import stock_code_map
from base.KeyTransform import trade_date_map

from gp_factor_construction.get_gp_factor_to_lib import GPFactorToLib

## TODO - 参数
start_date = '2017-01-01'
end_date = '2019-12-31'

original_db = 'gp_factors'
barra_stripping_db = 'gp_factors/barra_stripping'

# TODO - 因子选择的参数
long_group_excess_cond_threshold = 0.08
corr_threshold = 0.6

use_factor_test_api_cache = True

factor_test_config = {
    'start_date': start_date,
    'end_date':  end_date,
    'group_method': 'cap',
    # 'weight_method': 'cir_cap',
    'weight_method': 'ew',
    'group_quantile': 10,
    'neutral_group_quantile': 10,
    'rebalance_type': 'fixed',
    'rebalance_periods': 10,
    'universe_pool': 'all',
    'benchmark_index_code': '000905',
    'industry_type': 'ci1',
    'liquidity_filter': 'money',
    'liquidity_filter_period': 21,
    'factor_database': 'validation',
}

## TODO - 定义计算相关性的函数
def get_cross_sectional_corr(df, method='pearson', min_pct=0.5):
    df = df.copy()

    total_num = df.shape[0]

    corr_mat = df.corr(method, min_periods=int(total_num*min_pct))

    return corr_mat

def get_factor_corr(factor_data, method='pearson', min_pct=0.5):
    factor_data = factor_data.copy()

    corr_df = factor_data.groupby('date').apply(get_cross_sectional_corr,
                                                method=method,
                                                min_pct=min_pct)

    corr_df.index.names = ['date', 'factor_name']

    return corr_df

def get_mean_corr_data(corr_data, by='factor_name'):
    mean_corr_df = corr_data.groupby(by).mean()

    return mean_corr_df

#######
# TODO - API初始化
# factor_barra_strip_api = UsedFactorBarraStripping()

if not use_factor_test_api_cache:
    factor_test_api = FactorTest(
        start_date=factor_test_config['start_date'],
        end_date=factor_test_config['end_date'],
        group_quantile=factor_test_config['group_quantile'],
        neutral_group_quantile=factor_test_config['neutral_group_quantile'],
        rebalance_periods=factor_test_config['rebalance_periods'],
        rebalance_type=factor_test_config['rebalance_type'],
        universe_pool=factor_test_config['universe_pool'],
        benchmark_index_code=factor_test_config['benchmark_index_code'],
        liquidity_filter=factor_test_config['liquidity_filter'],
        liquidity_filter_period=factor_test_config['liquidity_filter_period'],
        group_method=factor_test_config['group_method'],
        weight_method=factor_test_config['weight_method'],
        industry_type=factor_test_config['industry_type'],
        factor_database=factor_test_config['factor_database'],
        # factor_lib_dir=factor_test_config['factor_lib_dir']
    )

    # TODO - 数据初始化
    factor_test_api.initialization()

    # cache api
    writeToFile(file=factor_test_api,
                store_path='./zbc_gplearn_factor_mining/buffer/factor_test_api_cache_10d',
                mode='wb')
else:
    factor_test_api = loadFile('./zbc_gplearn_factor_mining/buffer/factor_test_api_cache_10d')


def gp_factor_selection_v1(raw_data,
                           keys,
                           neu_keys,
                           model,
                           factor_name,
                           by_name=True,
                           key_by_path=False,
                           ts_base=5,
                           key_reverse=True):
    st = time()

    raw_data = raw_data.copy()

    # TODO - 传入模型，计算因子
    factor_data = get_factor_by_model(model=model,
                                      model_name=factor_name,
                                      start_date=None,
                                      end_date=None,
                                      raw_data=raw_data,
                                      keys=keys,
                                      neu_keys=neu_keys,
                                      by_name=by_name,
                                      key_by_path=key_by_path,
                                      ts_base=ts_base,
                                      key_reverse=key_reverse)

    print('get factor data done!')

    # TODO - 去掉nan的因子值
    factor_data = factor_data.dropna()

    # TODO - 进行barra剥离
    factor_data = factor_barra_strip_api.factor_barra_stripping(factor_data=factor_data, factor_name=factor_name)
    print('barra stripping done!')

    # TODO - 因子测试
    # 全市场
    long_group_excess, \
    long_group_excess_cond, \
    direction, \
    quantile_result_stats = factor_test_api.one_factor_analysis(factor_data=factor_data,
                                                                factor_name=factor_name,
                                                                fillna=False,
                                                                winsorize=False,
                                                                normalization=False,
                                                                drop_duplicates=True,
                                                                long_group_excess_cond_threshold=long_group_excess_cond_threshold)
    print('%s:\ndirection is %.0f and long group excess retun is %.2f%% sec.(is select: %s).\n' %
          (factor_name, direction, long_group_excess*100, long_group_excess_cond))

    # TODO - 因子相关性检验
    # 如果通过检验,则进行相关性分析
    if long_group_excess_cond:
        # TODO - 读取因子数据
        # data_reader = DataReader(db=original_db)
        data_reader = DataReader(db=barra_stripping_db)

        existing_factor_list = data_reader.show_factor_library_db()
        existing_factor_list = [ef.split('.h5')[0] for ef in existing_factor_list if '.h5' in ef]
        is_seleted = True

        if len(existing_factor_list) == 0:
            pass
        else:
            analysis_date_list = get_every_fixed_frequency_trade_date(start_date=start_date,
                                                                      end_date=end_date,
                                                                      freq=10)

            analysis_date_list = analysis_date_list['rebalance_start']

            analysis_date_list = analysis_date_list[(analysis_date_list >= start_date) &
                                                    (analysis_date_list <= end_date)]

            try:
                factor_data = factor_data.set_index(['date', 'stock_code'])
            except:
                pass

            for ef in existing_factor_list:
                # ef = existing_factor_list[0]
                efactor_data = data_reader.read_factor_table(ef)

                efactor_data = efactor_data[efactor_data['date'].isin(analysis_date_list)]

                efactor_data = efactor_data.set_index(['date', 'stock_code'])

                analysis_data = pd.concat([efactor_data, factor_data], axis=1,
                                          join_axes=[efactor_data.index])

                corr_data = get_factor_corr(analysis_data, method='pearson', min_pct=0.5)

                corr_data = get_mean_corr_data(corr_data, by='factor_name')

                corr_data = corr_data.iloc[0, 1]

                if corr_data > corr_threshold:
                    is_seleted = False

                    print('%s and %s correlation are %.4f > %.4f, so no selection.' %
                          (factor_name, ef, corr_data, corr_threshold))
                    break
    else:
        is_seleted = False

    # print('correlation filteration done!')
    et = time()
    print('%s selection - time spent is %.4f sec.\n' % (factor_name, et-st))

    return is_seleted



